#include "Planif.hh"
#include <vector>
#include <map>
#include <algorithm>
#include <limits>
#include "Bug.h"
#include "munkres.h"
#include <map>
#include "Bug.h"
#include <iostream>
#include <limits>


struct LocIndex : public std::map<uLocation,unsigned>
{
  std::vector<uLocation> reverse_index;

  LocIndex() : reverse_index(0) {}
  
  
  void insert(uLocation loc)
  {
    std::map<uLocation,unsigned>::iterator it = find(loc);
    if(it == end())
      {
	std::map<uLocation,unsigned>::insert(std::make_pair(loc,reverse_index.size()));
	reverse_index.push_back(loc);
      }
  }
	
  unsigned get_index(uLocation loc)
  {
    std::map<uLocation,unsigned>::iterator it = find(loc);
    return it->second;
  }

  uLocation get_location(unsigned index)
  {
    return reverse_index.at(index);
  }
 
};



bool by_increasing_cost(const Option & a, const Option & b)
{
  return (a.cost<b.cost);
}


bool by_decreasing_cost(const Option & a, const Option & b)
{
  return (a.cost>b.cost);
}



bool by_ant(const Option & a, const Option & b)
{
  return (a.from<b.from);
}



/// remove duplicate options and sum up their values/costs
void sum_up(std::vector<Option> & options)
{
  if(options.empty()) return;

  sort(options.begin(),options.end());
  
  std::vector<Option>::iterator it = options.begin();
  std::vector<Option> new_options(1,*it);
  ++it;
  while(it!=options.end())
    {
      if(it->from == new_options.back().from and it->to == new_options.back().to)
	new_options.back().cost += it-> cost;
      else
	new_options.push_back(*it);
      ++it;
    }

  options = new_options;
}


std::set<std::pair<uLocation,uLocation> > get_minimal_cost_plan(std::vector<Option> & options)
{
  return get_an_optimal_plan(options,true);
}

std::set<std::pair<uLocation,uLocation> > get_maximal_cost_plan(std::vector<Option> & options)
{
  return get_an_optimal_plan(options,false);
}


std::set<std::pair<uLocation,uLocation> > 
get_an_optimal_plan(std::vector<Option> & options, bool minimize_cost)
{
  /// remove duplicate options and sum up their values/costs
  sum_up(options);

  // add some random
  random_shuffle(options.begin(),options.end());

  // search unconstrained options
  std::map<uLocation,unsigned> conflicts;
  for(std::vector<Option>::const_iterator it = options.begin();
      it !=options.end(); ++it)
    conflicts[it->to]++;
  
  if(minimize_cost)
    stable_sort(options.begin(),options.end(), by_increasing_cost);
  else
    stable_sort(options.begin(),options.end(), by_decreasing_cost);

  // build matrix with contrained options only
  // other (unconstrained) options go directly to the plan with best cost
  LocIndex ressources;
  LocIndex tasks;
  FastSet<unsigned> & happy_ants(visited);
  happy_ants.reset(); // clear the set
  std::set<std::pair<uLocation,uLocation> > plan;

  
  for(std::vector<Option>::const_iterator it = options.begin();
      it !=options.end(); ++it)
    if(not happy_ants.count(it->from))
      {
	if(not ressources.count(it->from) and conflicts[it->to]==1) // no conflict with other ants
	  {
	    happy_ants.insert(it->from);
	    plan.insert(std::pair<uLocation,uLocation>(it->from,it->to));
	    bug << "adding: " << torus.to_str(it->from,it->to) << std::endl;
	  }
	else
	  { 
	    bug << "postponing: " << torus.to_str(it->from,it->to) << std::endl;
	    ressources.insert(it->from);
	    tasks.insert(it->to);
	  }
      }

  bug << "happy ants: " << happy_ants.size()
      << std::endl;

  if(ressources.empty())
    return plan;
    
 
  signed nrows = ressources.size();
  signed ncols = tasks.size();
  Matrix<float> matrix(nrows,ncols);

  // Initialize matrix with huge costs.
  for (signed row = 0 ; row != nrows ; row++ ) {
    for (signed col = 0 ; col != ncols ; col++ ) {
      matrix(row,col) = std::numeric_limits<float>::max();
    }
  }

  // fill the cost matrix
  for(std::vector<Option>::const_iterator it = options.begin();
      it !=options.end(); ++it)
    if(not happy_ants.count(it->from))
      {
	signed row = ressources.get_index(it->from);
	signed col = tasks.get_index(it->to);
	
	if(minimize_cost)
	  {
	    matrix(row,col) = it->cost;
	  }
	else
	  {
	    matrix(row,col) = -it->cost;
	  }
      }

#if DEBUG
  // Display cost matrix.
  bug << "cost matrix" << std::endl;
  for ( int row = 0 ; row < nrows ; row++ ) {
    for ( int col = 0 ; col < ncols ; col++ ) {
      bug.file.width(2);
      if(matrix(row,col) < std::numeric_limits<float>::max())
	bug << matrix(row,col) << ",";
      else
	bug << "inf" << ",";
    }
    bug << std::endl;
  }
#endif /* DEBUG */


  
  // Apply Munkres algorithm to matrix.
  bug << "optimize allocation" << std::endl;
  Munkres m;
  m.solve(matrix);

#if DEBUG
  // Display solved matrix.
  bug << "solved matrix" << std::fflush << std::endl;
  for ( int row = 0 ; row < nrows ; row++ ) {
    for ( int col = 0 ; col < ncols ; col++ ) {
      bug.file.width(2);
      bug << matrix(row,col)<< ",";
    }
    bug << std::endl;
  }
#endif /* DEBUG */


  for ( int row = 0 ; row != nrows ; row++ ) 
    {
      std::vector<unsigned> bestcol;
      int rowcount = 0;
      for ( int col = 0 ; col < ncols ; col++  ) 
	{
	  if ( matrix(row,col) == 0 )
	    {
	      rowcount++;
	      bestcol.push_back(col);
	    }
	}
      if ( rowcount != 1 )
	bug << "Row " << row << " has " << rowcount << " columns that have been matched." << std::endl;

      if(not bestcol.empty())
	{
	  signed select = bestcol.at(rand()%bestcol.size());
	  plan.insert(std::pair<uLocation,uLocation>(ressources.get_location(row),tasks.get_location(select)));
	  bug << "adding: " << torus.to_str(ressources.get_location(row),tasks.get_location(select)) << std::endl;
	}
      else
	{
	  bug << "No valid solution found for " << ressources.get_location(row) << std::endl;
	}
    }

  return plan;
}






